use async_trait::async_trait;
use base64::{engine::general_purpose::URL_SAFE_NO_PAD, Engine};
use hmac::{Hmac, Mac};
use reqwest::Client;
use serde_json::{json, Value};
use sha2::Sha256;
use std::time::{SystemTime, UNIX_EPOCH};
use futures::StreamExt;

use crate::error::{Result, ConversionError};
use crate::config::settings::{SYSTEM_PROMPT, ApiSettings};

#[async_trait]
pub trait ApiClient: Send + Sync {
    async fn generate_code(&self, code: &str, prompt: &str) -> Result<String>;
}

pub struct ChatGLMClient {
    settings: ApiSettings,
    client: Client,
    jwt_token: String,
}

impl ChatGLMClient {
    pub fn new(settings: ApiSettings) -> Result<Self> {
        let client = reqwest::Client::builder()
            .timeout(std::time::Duration::from_secs(settings.timeout))
            .build()?;

        let jwt_token = Self::generate_jwt(&settings.key);

        Ok(Self {
            settings,
            client,
            jwt_token,
        })
    }

    fn generate_jwt(api_key: &str) -> String {
        let (user_id, user_secret) = api_key.split_once('.').unwrap();
        
        let header = json!({
            "alg": "HS256",
            "sign_type": "SIGN"
        });
        
        let timestamp = SystemTime::now()
            .duration_since(UNIX_EPOCH)
            .unwrap()
            .as_secs();

        let payload = json!({
            "api_key": user_id,
            "exp": timestamp * 3,
            "timestamp": timestamp
        });

        let encoded_header = URL_SAFE_NO_PAD.encode(header.to_string().as_bytes());
        let encoded_payload = URL_SAFE_NO_PAD.encode(payload.to_string().as_bytes());
        
        let signature_input = format!("{}.{}", encoded_header, encoded_payload);
        
        type HmacSha256 = Hmac<Sha256>;
        let mut mac = HmacSha256::new_from_slice(user_secret.as_bytes()).unwrap();
        mac.update(signature_input.as_bytes());
        let signature = mac.finalize().into_bytes();
        
        let encoded_signature = URL_SAFE_NO_PAD.encode(&signature);
        
        format!("{}.{}", signature_input, encoded_signature)
    }
}

#[async_trait]
impl ApiClient for ChatGLMClient {
    async fn generate_code(&self, code: &str, prompt: &str) -> Result<String> {
        let prompt = prompt.replace("{code}", code);
        
        let data = json!({
            "model": "glm-4",
            "messages": [
                {"role": "system", "content": SYSTEM_PROMPT},
                {"role": "user", "content": prompt}
            ],
            "stream": true
        });

        let response = self.client
            .post(&self.settings.base_url)
            .header("Authorization", format!("Bearer {}", self.jwt_token))
            .header("Content-Type", "application/json")
            .json(&data)
            .send()
            .await?;

        if !response.status().is_success() {
            return Err(ConversionError::Conversion(
                format!("API request failed with status: {}", response.status())
            ));
        }

        let mut full_response = String::new();
        let mut stream = response.bytes_stream();

        while let Some(chunk) = stream.next().await {
            let chunk = chunk?;
            if let Ok(line) = String::from_utf8(chunk.to_vec()) {
                if line.starts_with("data: ") {
                    let json_str = &line[6..];
                    if json_str != "[DONE]" {
                        if let Ok(value) = serde_json::from_str::<Value>(json_str) {
                            if let Some(content) = value["choices"][0]["delta"]["content"].as_str() {
                                full_response.push_str(content);
                            }
                        }
                    }
                }
            }
        }

        if full_response.trim().is_empty() {
            return Err(ConversionError::Conversion("Empty response from API".to_string()));
        }

        Ok(full_response)
    }
} 